FSDP Colocate Mode OOM Issues in verl-omni
The problem is that when configuring student and teacher model to colocate on same GPU/NPU, building the teacher model encounters OOM issues. This post introduces a simple fix to this problem.
Table of Contents
1. Problem Statement
When configuring student and teacher model to colocate on the same GPU/NPU (as in DiffusionOPD snippets), teachers’ FSDP build might fail due to OOM after students’ FSDP fail.
2. Possible Fix
A possible fix is that, to add the following argument:
+actor_rollout_ref.ref.fsdp_config.wrap_policy.transformer_layer_cls_to_wrap=[LTX2VideoTransformerBlock] # replace the LTX2VideoTransformerBlock with classnames
3. Investigation
This Hydra configuration tells FSDP’s auto-wrap policy to make each submodule whose class is LTX2VideoTransformerBlock its own FSDP sharding unit, instead of wrapping the entire transformer as one single FSDP sharding unit.
3.1. Mechanism
The FSDP (torch’s FullyShardedDataParallel) requires a “wrap policy” to determine the sharding boundries. When it wraps the whole model as a single unit, it calls _init_param_handle_from_params() to flatten all parameters of that unit into one contiguous flat buffer via flatten_tensors_into_flat_param(), which is a simple torch.cat(). Therefore, in my case, a 17.7B LTX2.3 in BF16 requires 35GB of HBM, causing the OOM.
By setting the transformer_layer_cls_to_wrap, each LTX2VideoTransformerBlock has its own FSDP unit. Now each unit is only around 0.7GB and is loaded on-demand.